# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from quant_researcher.quant.datasource_fetch.index_api import index_components
from quant_researcher.quant.factors.factor_preprocess.preprocess import my_qcut2, ts_percentile_rank_score
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.project_tool.common_var import PY_DATE_FMT_WITHOUT_N_DASH, PY_DATE_FMT_WITH_N_DASH
from quant_researcher.quant.project_tool.wrapper_tools.common_wrappers import count_time

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

bp_table = 'stk_fac_fdmt_value_bp'
ep_table = 'stk_fac_fdmt_value_etop'
sp_table = 'stk_fac_fdmt_value_ps'
turnover_table = 'stk_fac_hq_liquidity_vol63'
beta_table = 'stk_fac_risk_beta63'
#
# def reset_connection():
#     print('数据库连接重建')
#     global CON, CON_factor, CON_astk, CON_ty, CON_wst
#     CON.close()
#     CON_factor.close()
#     CON_astk.close()
#     CON_ty.close()
#     CON_wst.close()
#     # 基础数据组数据库
#     CON = db_conn.create_a_factor_db_conn()
#     # 投研组衍生指标数据库
#     CON_ty = db_conn.get_derivative_factor_db_conn()
#     # 因子数据库
#     CON_factor = db_conn.get_tk_factors_conn()
#     # 维赛特数据库
#     CON_wst = db_conn.get_longbo_conn()
#     CON_astk = db_conn.get_a_stk_conn()


class FactorCrowding(object):
    def __init__(self, factor_name='',
                 factor_direction=1,
                 group_nums=5,
                 begin_date='',
                 end_date='',
                 universe='000985',
                 external_data=None,
                 quantile=True):
        # 基础数据组数据库
        self.con = db_conn.get_basic_data_conn(by_sqlalchemy=True)
        # 基础股票指标数据库
        self.con_astk = db_conn.get_stock_conn(by_sqlalchemy=True)
        # 因子数据库
        self.con_factor = db_conn.get_tk_factors_conn(by_sqlalchemy=True)
        # 维赛特数据库
        self.begin_date = begin_date
        self.end_date = end_date
        self.factor_name = factor_name
        self.factor_direction = factor_direction
        self.group_nums = group_nums
        self.external_data = external_data
        self.universe = universe
        self.quantile = quantile
        self.stock_pool = index_components.get_index_hist_stock_pool(self.universe,
                                                                     self.begin_date, self.end_date)['stock_code'].tolist()
        if external_data is not None:
            self.factor_table_name = None
            self.data_begin_date = self.begin_date
        else:
            self.factor_table_name, self.factor_name_cn = self.get_factor_table_name()
            self.data_begin_date = pd.read_sql(f"select distinct tradedate from {self.factor_table_name} "
                                               f"where tradedate < '{self.begin_date}' order by tradedate desc "
                                               f"limit 63", self.con_factor).iloc[-1, 0].strftime(PY_DATE_FMT_WITHOUT_N_DASH)
        self.data_end_date = self.end_date.replace('-', '')
        self.factor_data = self.get_factor_data()
        self.mask_top, self.mask_bottom = self.cut_group()
        self.stock_ret = self.get_stock_ret()
        self.bp, self.ep, self.sp = self.get_value_spread_data()
        self.turnover = self.get_turnover()
        self.beta = self.get_beta()

    def get_factor_table_name(self):
        factor_table_name = pd.read_sql("select factor_table_name, factor_name_cn from stk_fac_info "
                                             f"where factor_name_en = '{self.factor_name}'", self.con_factor)
        factor_name_cn = factor_table_name.iloc[0, 1]
        return factor_table_name.iloc[0, 0], factor_name_cn

    def update_data(self, data=None, factor_name=None, direction=1, quantile=True):
        if data is not None:
            self.external_data = data
            self.factor_data = self.get_factor_data()
        else:
            assert factor_name is not None, '请至少输入一个新的因子数据或名称'
            if self.factor_name == factor_name:
                print('已经下载所需数据')
                pass
            else:
                self.factor_name = factor_name
                self.factor_table_name, self.factor_name_cn = self.get_factor_table_name()
                print('获取新因子数据')
                self.factor_data = self.get_factor_data()
        self.factor_direction = direction
        self.mask_top, self.mask_bottom = self.cut_group(quantile=quantile)

    @count_time
    def get_factor_data(self):
        if self.external_data is not None:
            factor_unstack = self.external_data
        else:
            factor_data = pd.read_sql(f"select tradedate, stockcode, value as {self.factor_name} "
                                      f"from {self.factor_table_name} "
                                      f"where tradedate >= '{self.begin_date}' "
                                      f"and tradedate <= '{self.end_date}' "
                                      f"and stockcode in {tuple(self.stock_pool)} "
                                      f"order by tradedate ", self.con_factor)
            factor_data.dropna(inplace=True)
            factor_data[self.factor_name] = factor_data[self.factor_name].astype(float)
            factor_data['tradedate'] = factor_data['tradedate'].astype('str')
            factor_unstack = factor_data.set_index(['tradedate', 'stockcode'])[self.factor_name].unstack()
        return factor_unstack

    @count_time
    def cut_group(self, quantile=True):
        assert 1 < self.group_nums <= 10, "分组数应大于1且不大于10"
        assert isinstance(self.group_nums, int), "分组数必须为整数"
        if quantile:
            lower = 1 / self.group_nums
            upper = 1 - lower
            high = self.factor_data.quantile(q=upper, axis=1)
            low = self.factor_data.quantile(q=lower, axis=1)
            higher_part = self.factor_data.gt(high, axis=0)
            lower_part = self.factor_data.lt(low, axis=0)
            mask_top = higher_part if self.factor_direction == 1 else lower_part
            mask_bottom = higher_part if self.factor_direction == -1 else lower_part
        else:
            cut_group = self.factor_data.apply(my_qcut2, group_num=self.group_nums, axis=1, cut_type='non_quantile')
            higher_part = cut_group==self.group_nums
            lower_part = cut_group==1
            mask_top = higher_part if self.factor_direction == 1 else lower_part
            mask_bottom = higher_part if self.factor_direction == -1 else lower_part
        return mask_top, mask_bottom

    @count_time
    def get_stock_ret(self):
        close = pd.read_sql("select sec_code as stockcode, trade_date as tradedate, "
                            "close as fap_close "
                            "from astk_qt_stockdailyquote_adjbf "
                            f"where trade_date >='{self.data_begin_date}' "
                            f"and trade_date <='{self.data_end_date}' "
                            f"and sec_code in {tuple(self.stock_pool)} ", self.con_astk)
        close = close[close.stockcode.str[0].isin(['0', '3', '6'])]
        close['tradedate'] = pd.to_datetime(close['tradedate']).dt.strftime(PY_DATE_FMT_WITH_N_DASH)
        ret = close.set_index(['tradedate', 'stockcode'])['fap_close'].unstack().ffill().pct_change()
        return ret

    @count_time
    def get_value_spread_data(self):
        bp = pd.read_sql(f"select tradedate, stockcode, value as bp "
                         f"from {bp_table} "
                         f"where tradedate >= '{self.begin_date}' "
                         f"and tradedate <= '{self.end_date}' "
                         f"and stockcode in {tuple(self.stock_pool)} ",
                         self.con_factor, index_col=['tradedate', 'stockcode'])['bp'].unstack()
        ep = pd.read_sql(f"select tradedate, stockcode, value as ep "
                         f"from {ep_table} "
                         f"where tradedate >= '{self.begin_date}' "
                         f"and tradedate <= '{self.end_date}' "
                         f"and stockcode in {tuple(self.stock_pool)} "
                         , self.con_factor, index_col=['tradedate', 'stockcode'])['ep'].unstack()
        sp = 1 / pd.read_sql(f"select tradedate, stockcode, value as sp "
                             f"from {sp_table} "
                             f"where tradedate >= '{self.begin_date}' "
                             f"and tradedate <= '{self.end_date}' "
                             f"and stockcode in {tuple(self.stock_pool)} "
                             , self.con_factor, index_col=['tradedate', 'stockcode'])['sp'].unstack()
        return bp, ep, sp

    @count_time
    def get_turnover(self):
        turnover = pd.read_sql(f"select tradedate, stockcode, value as turnover "
                               f"from {turnover_table} "
                               f"where tradedate >= '{self.begin_date}' "
                               f"and tradedate <= '{self.end_date}' "
                               f"and stockcode in {tuple(self.stock_pool)} ",
                               self.con_factor, index_col=['tradedate', 'stockcode'])['turnover'].unstack()
        return turnover

    @count_time
    def get_beta(self):
        beta = pd.read_sql(f"select tradedate, stockcode, value as beta "
                           f"from {beta_table} "
                           f"where tradedate >= '{self.begin_date}' "
                           f"and tradedate <= '{self.end_date}' "
                           f"and stockcode in {tuple(self.stock_pool)} "
                           , self.con_factor, index_col=['tradedate', 'stockcode'])['beta'].unstack()
        return beta

    @count_time
    def get_value_spread(self, plot=True):  # 估值价差
        bp, ep, sp = self.bp, self.ep, self.sp
        bp_spread = np.log(bp[self.mask_top].median(axis=1) / bp[self.mask_bottom].median(axis=1))
        sp_spread = np.log(sp[self.mask_top].median(axis=1) / sp[self.mask_bottom].median(axis=1))
        ep_spread = ep[self.mask_top].median(axis=1) - ep[self.mask_bottom].median(axis=1)
        # 以200日窗口进行标准化，而非全部历史数据
        bp_spread = (bp_spread - bp_spread.rolling(504, min_periods=480).mean()) / bp_spread.rolling(504, min_periods=480).std()
        sp_spread = (sp_spread - sp_spread.rolling(504, min_periods=480).mean()) / sp_spread.rolling(504, min_periods=480).std()
        ep_spread = (ep_spread - ep_spread.rolling(504, min_periods=480).mean()) / ep_spread.rolling(504, min_periods=480).std()
        result = (bp_spread + sp_spread + ep_spread) / 3
        result = result.dropna()
        result.name = 'value_spread'
        self.value_spread = result
        if plot:
            self.plot(self.value_spread)

    @count_time
    # 配对相关性
    def get_pair_corr(self, plot=True):
        ret = self.stock_ret
        ex_mean = ((-ret).add(ret.sum(axis=1), axis=0)).div(ret.transform('count', axis=1) - 1, axis=0)  # 剔除自己后的组均值
        correlation = ret.rolling(63).corr(ex_mean).replace([np.inf, -np.inf], 0)
        top_corr = correlation[self.mask_top].mean(axis=1)
        bottom_corr = correlation[self.mask_bottom].mean(axis=1)
        avg_corr = (top_corr + bottom_corr) / 2
        # 在综合函数部分标准化，计算内部不进行标准化
        # avg_corr = (avg_corr - avg_corr.expanding().mean()) / avg_corr.expanding().std()
        avg_corr = avg_corr.dropna()
        avg_corr.name = 'pair_corr'
        avg_corr.index = pd.to_datetime(avg_corr.index)
        self.pair_corr = avg_corr.loc[self.begin_date:]
        if plot:
            self.plot(self.pair_corr)

    @count_time
    # 波动率比率
    def get_vol_ratio(self, plot=True):
        ret = self.stock_ret
        volatility = ret.rolling(63).std()
        volatility_ratio = volatility[self.mask_top].median(axis=1) / volatility[self.mask_bottom].median(axis=1)
        # 在综合函数部分标准化，计算内部不进行标准化
        # volatility_ratio = (volatility_ratio - volatility_ratio.expanding().mean()) / volatility_ratio.expanding().std()
        volatility_ratio = volatility_ratio.dropna()
        volatility_ratio.name = 'vol_ratio'
        volatility_ratio.index = pd.to_datetime(volatility_ratio.index)
        self.vol_ratio = volatility_ratio.loc[self.begin_date:]
        if plot:
            self.plot(self.vol_ratio)

    @count_time
    # 换手率比率
    def get_turnover_ratio(self, plot=True):
        vol = self.turnover
        vol_ratio = vol[self.mask_top].median(axis=1) / vol[self.mask_bottom].median(axis=1)
        # 在综合函数部分标准化，计算内部不进行标准化
        # turn_ratio = (vol_ratio - vol_ratio.expanding().mean()) / vol_ratio.expanding().std()
        turn_ratio = vol_ratio.dropna()
        turn_ratio.name = 'turnover_ratio'
        self.turnover_ratio = turn_ratio
        if plot:
            self.plot(self.turnover_ratio)

    @count_time
    # beta比率
    def get_beta_ratio(self, plot=True):
        beta = self.beta
        beta_ratio = beta[self.mask_bottom].median(axis=1) / beta[self.mask_top].median(axis=1)
        # 在综合函数部分标准化，计算内部不进行标准化
        # beta_ratio = (beta_ratio - beta_ratio.expanding().mean()) / beta_ratio.expanding().std()
        beta_ratio = beta_ratio.dropna()
        beta_ratio.name = 'beta_ratio'
        self.beta_ratio = beta_ratio
        if plot:
            self.plot(self.beta_ratio)

    @count_time
    def get_composite(self, raw=False, smooth=True, standardize=True, expanding_rank=False, plot=False):
        self.get_value_spread(plot=plot)
        self.get_turnover_ratio(plot=plot)
        self.get_pair_corr(plot=plot)
        self.get_vol_ratio(plot=plot)
        self.get_beta_ratio(plot=plot)
        if smooth:
            self.turnover_ratio = self.turnover_ratio.rolling(63).mean()
            self.value_spread = self.value_spread.rolling(63).mean()
            self.beta_ratio = self.beta_ratio.rolling(63).mean()
            self.vol_ratio = self.vol_ratio.rolling(63).mean()
            self.pair_corr = self.pair_corr.rolling(63).mean()
        if standardize:
            # value_spread 由于合成算法的原因，已经在计算时进行了标准化，所以不再标准化
            self.turnover_ratio = (self.turnover_ratio - self.turnover_ratio.rolling(504, min_periods=480).mean()) / self.turnover_ratio.rolling(504, min_periods=480).std()
            self.pair_corr = (self.pair_corr - self.pair_corr.rolling(504, min_periods=480).mean()) / self.pair_corr.rolling(504, min_periods=480).std()
            self.vol_ratio = (self.vol_ratio - self.vol_ratio.rolling(504, min_periods=480).mean()) / self.vol_ratio.rolling(504, min_periods=480).std()
            self.beta_ratio = (self.beta_ratio - self.beta_ratio.rolling(504, min_periods=480).mean()) / self.beta_ratio.rolling(504, min_periods=480).std()
        all_df = pd.concat([self.value_spread, self.beta_ratio, self.turnover_ratio, self.vol_ratio, self.pair_corr], axis=1)
        if expanding_rank:
            all_df = all_df.apply(ts_percentile_rank_score, way='rolling', min_period=504, scale=100)
            # self.value_spread = ts_rank(self.value_spread, way='expanding', min_period=126, scale=100).dropna()
            # self.pair_corr = ts_rank(self.pair_corr, way='expanding', min_period=126, scale=100).dropna()
            # self.vol_ratio = ts_rank(self.vol_ratio, way='expanding', min_period=126, scale=100).dropna()
            # self.turnover_ratio = ts_rank(self.turnover_ratio, way='expanding', min_period=126, scale=100).dropna()
            # self.beta_ratio = ts_rank(self.beta_ratio, way='expanding', min_period=126, scale=100).dropna()
        if not raw:
            composite = all_df.mean(axis=1)
            return composite
        else:
            return all_df

    def plot_all(self):
        all_df = pd.concat([self.value_spread, self.beta_ratio, self.turnover_ratio, self.vol_ratio, self.pair_corr],
                           axis=1)
        all_df['综合拥挤度'] = all_df.mean(axis=1)
        tick = all_df.index.strftime('%Y-%m')
        label = [tick[i] if i % 21 == 0 else '' for i in range(len(tick))]
        ax = all_df.plot.bar(figsize=(22, 18), sharex=False, subplots=True, layout=(3, 2), legend=False)
        # fig, ax = plt.subplots(3, 2, figsize=(22, 18))
        # ax[0, 0].bar(self.value_spread.index, self.value_spread)
        # ax[0, 1].bar(self.pair_corr.index, self.pair_corr)
        # ax[1, 0].bar(self.vol_ratio.index, self.vol_ratio)
        # ax[1, 1].bar(self.turnover_ratio.index, self.turnover_ratio, 3)
        # ax[2, 0].bar(self.beta_ratio.index, self.beta_ratio, 3)
        # ax[2, 1].bar(avg_crowding.index, avg_crowding, 3)
        for x, method in zip(ax.reshape(-1), all_df.columns):
            x.set_title(f'{method}')
            x.xaxis.set_major_formatter(ticker.FixedFormatter(label))
        plt.suptitle(f'{self.factor_name} 因子拥挤度', fontsize=18, y=0.92)
        plt.tight_layout(rect=[0, 0, 1, 0.9])
        plt.show()
        # plt.show()
        # plt.savefig('因子拥挤度_%s.png' % end)

    def plot(self, method_series):
        ax = method_series.plot.bar(figsize=(16, 9), color='g')
        # fig, ax = plt.subplots(figsize=(16, 9))
        tick = method_series.index.strftime('%Y-%m')
        label = [tick[i] if i % 21 == 0 else '' for i in range(len(tick))]
        ax.xaxis.set_major_formatter(ticker.FixedFormatter(label))
        # ax.bar(range(len(method_series)), method_series)
        ax.set_title(f'{method_series.name}')
        plt.show()


# %%
if __name__ == '__main__':
    factor = FactorCrowding('etop', 1, 5, '2018-01-01', '2019-06-30', universe='000300')
    # factor.get_value_spread(plot=True)
    # factor.get_pair_corr()
    # factor.vol_ratio()
    # factor.turnover_ratio()
    # factor.beta_ratio()
    test = factor.get_composite(raw=True, standardize=True, expanding_rank=True, plot=False)
